#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
数据处理模块
负责准备训练数据集，支持经验回放防遗忘
"""

import json
import os
from datasets import Dataset


def prepare_dataset(data_path, tokenizer, max_length=512, old_data_path=None, replay_ratio=0.0):
    """
    准备训练数据集，支持经验回放防遗忘
    
    Args:
        data_path: 新数据路径
        tokenizer: 分词器
        max_length: 最大长度
        old_data_path: 旧数据路径（用于经验回放）
        replay_ratio: 回放比例
    """
    print(f"正在加载数据集: {data_path}")
    
    # 读取新数据
    with open(data_path, 'r', encoding='utf-8') as f:
        new_data = [json.loads(line) for line in f]
    
    # 经验回放：混入旧数据
    if old_data_path and replay_ratio > 0 and os.path.exists(old_data_path):
        with open(old_data_path, 'r', encoding='utf-8') as f:
            old_data = [json.loads(line) for line in f]
        
        replay_count = int(len(new_data) * replay_ratio)
        replay_data = old_data[:min(replay_count, len(old_data))]
        
        training_data = new_data + replay_data
        print(f"🔄 经验回放: {len(new_data)}新 + {len(replay_data)}旧 = {len(training_data)}总")
    else:
        training_data = new_data
        print(f"📊 训练数据: {len(training_data)}条")
    
    # 格式化数据
    def format_example(example):
        prompt = example["prompt"]
        response = example["response"]
        full_text = prompt + response
        
        # 编码
        encoding = tokenizer(
            full_text,
            truncation=True,
            max_length=max_length,
            padding="max_length",
            return_tensors=None,
        )
        
        # 创建正确的labels，将prompt部分设为-100
        labels = encoding["input_ids"].copy()
        
        # 计算prompt的长度
        prompt_encoding = tokenizer(
            prompt,
            truncation=True,
            max_length=max_length,
            add_special_tokens=False
        )
        prompt_length = len(prompt_encoding["input_ids"])
        
        # 将prompt部分的label设为-100
        labels[:prompt_length] = [-100] * prompt_length
        
        return {
            "input_ids": encoding["input_ids"],
            "attention_mask": encoding["attention_mask"],
            "labels": labels
        }
    
    # 转换为Dataset对象
    dataset = Dataset.from_list(training_data)
    print("正在处理数据集...")
    dataset = dataset.map(format_example, remove_columns=dataset.column_names)
    print("数据集处理完成")
    
    return dataset


if __name__ == "__main__":
    # 可以添加一些测试代码
    pass